import numpy as np
import matplotlib.pyplot as plt


class LaneChangeReferenceGenerator:
    def __init__(self, config=None):
        # 参数配置
        self.default_config = {
            'dt': 0.1,  # 时间分辨率（秒）
            'predict_time': 3.0,  # 变道总时间（秒）
            'max_accel': 1.0,  # 最大横向加速度（m/s²）
            'min_speed': 1.0,  # 最小纵向速度（m/s）
            'max_speed': 10.0  # 最大纵向速度（m/s）
        }
        self.config = config or self.default_config

    def generate_reference(self, current_state, target_y):
        """
        生成变道参考线
        :param current_state: [x, y, v, heading] 当前状态（m, m, m/s, rad）
        :param target_y: 目标车道的Y坐标（m）
        :return: 参考线轨迹 [[x0,y0], [x1,y1], ...]
        """
        x0, y0, v, psi = current_state
        dy = target_y - y0

        # 纵向（X方向）运动参数
        T = self.config['predict_time']
        t_samples = np.arange(0, T + self.config['dt'], self.config['dt'])

        # 生成X坐标（匀速模型）
        x_traj = x0 + v * np.cos(psi) * t_samples

        # 生成Y坐标（五次多项式）
        y_traj = self._quintic_polynomial_planning(
            y0, dy,
            t_start=0, t_end=T,
            dy_start=0, dy_end=0,  # 初始和终止横向速度
            ddy_start=0, ddy_end=0  # 初始和终止横向加速度
        )

        # 截断到预测时间
        ref_traj = np.column_stack((x_traj[:len(y_traj)], y_traj))

        # 可视化
        # self._visualize(ref_traj, current_state, target_y)
        return ref_traj

    def _quintic_polynomial_planning(self, y0, dy, t_start, t_end,
                                     dy_start, dy_end, ddy_start, ddy_end):
        """
        五次多项式轨迹生成（Y方向）
        :param y0: 初始Y坐标
        :param dy: Y方向总位移
        :param t_start: 起始时间（通常为0）
        :param t_end: 终止时间
        :param dy_start: 初始横向速度
        :param dy_end: 终止横向速度
        :param ddy_start: 初始横向加速度
        :param ddy_end: 终止横向加速度
        """
        A = np.array([
            [t_end ** 3, t_end ** 4, t_end ** 5],
            [3 * t_end ** 2, 4 * t_end ** 3, 5 * t_end ** 4],
            [6 * t_end, 12 * t_end ** 2, 20 * t_end ** 3]
        ])

        b = np.array([
            dy - (dy_start * t_end + 0.5 * ddy_start * t_end ** 2),
            dy_end - (dy_start + ddy_start * t_end),
            ddy_end - ddy_start
        ])

        try:
            a = np.linalg.solve(A, b)
        except np.linalg.LinAlgError:
            return np.full(len(np.arange(t_start, t_end, self.config['dt'])), y0)

        # 时间序列
        t = np.arange(t_start, t_end, self.config['dt'])

        # 位置计算
        y = y0 + dy_start * t + 0.5 * ddy_start * t ** 2 + a[0] * t ** 3 + a[1] * t ** 4 + a[2] * t ** 5
        return y

    def _visualize(self, trajectory, current_state, target_y):
        """可视化结果"""
        plt.figure(figsize=(10, 6))

        # 绘制参考线
        plt.plot(trajectory[:, 0], trajectory[:, 1], 'b-', lw=2, label='生成轨迹')
        plt.axhline(y=current_state[1], color='g', linestyle='--', label='原始车道')
        plt.axhline(y=target_y, color='r', linestyle='--', label='目标车道')

        # 绘制车辆状态
        plt.quiver(current_state[0], current_state[1],
                   np.cos(current_state[3]), np.sin(current_state[3]),
                   scale=15, color='purple', label='初始状态')

        plt.grid(True)
        plt.axis('equal')
        plt.xlabel('纵向位置 X (m)')
        plt.ylabel('横向位置 Y (m)')
        plt.legend()
        plt.title(f'变道参考线生成（ΔY = {target_y - current_state[1]:.2f}m）')
        plt.show()


# 使用示例
if __name__ == "__main__":
    # 初始化生成器
    generator = LaneChangeReferenceGenerator()

    # 初始状态：车辆在y=1.5m的车道，以3m/s速度沿X轴行驶
    initial_state = [46, 1.7*3, 11.1, 0.0]

    # 目标车道y=4.5m
    target_lane_y = 1.7

    # 生成参考线
    ref_traj = generator.generate_reference(initial_state, target_lane_y)
    ref_traj = np.array(ref_traj)
    # 打印轨迹信息
    print("生成轨迹示例：")
    print(f"起点: {ref_traj[0]}")
    print(f"终点: {ref_traj[-1]}")
    print(type(ref_traj))
    print(ref_traj)
    print(f"轨迹点数: {len(ref_traj)}")